/*
* Created by F1shl3gs on 16/7/22.
*/

#include <ngx_config.h>
#include <ngx_core.h>
#include <ngx_http.h>
#include <ngx_string.h>


#define NGX_HTTP_ABTEST_SIZE        (8 * 1024)
#define NGX_HTTP_ABTEST_MASK        (NGX_HTTP_ABTEST_SIZE - 1)
#define NGX_HTTP_ABTEST_SHUFFIX     "F1shl3gs"

#define NGX_HTTP_ABTEST_PERCENT     1
#define NGX_HTTP_ABTEST_MATCH       2

typedef ngx_int_t (*ngx_http_abtest_handler_pt)(ngx_http_request_t *r, ngx_str_t *to);

typedef struct {
    uint32_t                        from;
    uint32_t                        to;

    ngx_str_t                       url;
} ngx_http_abtest_percent_part_t;

typedef struct {
    ngx_str_t                       pattern;
    ngx_regex_t                     regex;
    ngx_str_t                       url;
} ngx_http_abtest_match_part_t;


typedef struct {
    ngx_array_t                     parts;
    ngx_str_t                       key;
    ngx_http_abtest_handler_pt      handler;
    ngx_uint_t                      def;
} ngx_http_abtest_loc_conf_t;


/* static ngx_str_t                    prefix; */

static char *ngx_http_abtest_block(ngx_conf_t *cf, ngx_command_t *cmd, void *conf);
static char *ngx_http_abtest_match(ngx_conf_t *cf, ngx_command_t *cmd, void *conf);
static char *ngx_http_abtest_percent(ngx_conf_t *cf, ngx_command_t *cmd, void *conf);

static void *ngx_http_abtest_create_loc_conf(ngx_conf_t *cf);
static ngx_int_t ngx_http_abtest_postconfiguration(ngx_conf_t *cf);

static char *ngx_http_abtest_percent_init(ngx_conf_t *cf, ngx_http_abtest_loc_conf_t *alcf);
static char *ngx_http_abtest_match_init(ngx_conf_t *cf, ngx_http_abtest_loc_conf_t *alcf);
static ngx_int_t ngx_http_abtest_handler(ngx_http_request_t *r);


static ngx_int_t ngx_http_abtest_percent_handler(ngx_http_request_t *r, ngx_str_t *to);
static ngx_int_t ngx_http_abtest_match_handler(ngx_http_request_t *r, ngx_str_t *to);

static ngx_command_t ngx_http_abtest_commands[] = {
        {
                ngx_string("abtest"),
                NGX_HTTP_LOC_CONF | NGX_CONF_TAKE2 | NGX_CONF_BLOCK,
                ngx_http_abtest_block,
                NGX_HTTP_LOC_CONF_OFFSET,
                0,
                NULL
        },

        ngx_null_command
};

static ngx_http_module_t ngx_http_abtest_module_ctx = {
        NULL,                                  /* preconfiguration */
        ngx_http_abtest_postconfiguration,     /* postconfiguration */

        NULL,                                  /* create main configuration */
        NULL,                                  /* init main configuration */

        NULL,                                  /* create server configuration */
        NULL,                                  /* merge server configuration */

        ngx_http_abtest_create_loc_conf,       /* create location configuration */
        NULL                                   /* merge location configuration */
};

ngx_module_t    ngx_http_abtest_module = {
        NGX_MODULE_V1,
        &ngx_http_abtest_module_ctx,           /* module context */
        ngx_http_abtest_commands,              /* module commands */
        NGX_HTTP_MODULE,                       /* module type */
        NULL,                                  /* init master */
        NULL,                                  /* init module */
        NULL,                                  /* init process */
        NULL,                                  /* init thread */
        NULL,                                  /* exit thread */
        NULL,                                  /* exit process */
        NULL,                                  /* exit master */
        NGX_MODULE_V1_PADDING
};


static char *ngx_http_abtest_block(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) {
    ngx_conf_t                          save;
    ngx_uint_t                          i, mode;
    ngx_str_t                           *value;
    char                                *rv;
    ngx_http_abtest_loc_conf_t          *alcf;
    size_t                              as;

    alcf = conf;
    as = 0;
    mode = 0;
    value = cf->args->elts;

    for(i = 1; i < 3; ++i) {
        if(value[i].len == 5) {
            if(ngx_strncmp(value[i].data, "match", 5) == 0) {
                mode = NGX_HTTP_ABTEST_MATCH;
                as = sizeof(ngx_http_abtest_match_part_t);
                continue;
            }
        }

        if(value[i].len == 7) {
            if(ngx_strncmp(value[i].data, "percent", 7) == 0) {
                mode = NGX_HTTP_ABTEST_PERCENT;
                as = sizeof(ngx_http_abtest_percent_part_t);
                continue;
            }
        }

        if(value[i].len < 6) {
            goto invalid;
            return NGX_CONF_ERROR;
        }

        if(ngx_strncmp(value[i].data, "key=", 4) == 0) {
            alcf->key.data = value[i].data + 4;
            alcf->key.len = value[i].len - 4;
        }

    }

    if(ngx_array_init(&alcf->parts, cf->pool, 4, as) != NGX_OK) {
        return NGX_CONF_ERROR;
    }

    /* todo add shuffix support */

    save = *cf;
    cf->ctx = alcf;
    cf->handler = (mode == NGX_HTTP_ABTEST_MATCH)
                  ? ngx_http_abtest_match : ngx_http_abtest_percent;

    rv = ngx_conf_parse(cf, NULL);
    if(rv != NGX_CONF_OK) {
        return NGX_CONF_ERROR;
    }

    *cf = save;

    if(mode == NGX_HTTP_ABTEST_PERCENT) {
        rv = ngx_http_abtest_percent_init(cf, alcf);
    } else if(mode == NGX_HTTP_ABTEST_MATCH) {
        rv = ngx_http_abtest_match_init(cf, alcf);
    }

    return rv;

invalid:
    ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                       "invalid params \"%V\" for abtest", &value[i]);
    return NGX_CONF_ERROR;
}


static char *ngx_http_abtest_percent(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) {
    ngx_str_t                           *value;
    ngx_http_abtest_loc_conf_t          *alcf;
    ngx_http_abtest_percent_part_t      *part;
    ngx_int_t                           n;

    if(cf->args->nelts < 2 || cf->args->nelts > 3) {
        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                            "invalid param number");

        return NGX_CONF_ERROR;
    }

    alcf = cf->ctx;
    value = cf->args->elts;

    n = ngx_atoi(value->data, value->len);
    if(n == 0 || n == NGX_ERROR || n >= 100) {
        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                            "invalid param", &value[0]);

        return NGX_CONF_ERROR;
    }

    part = ngx_array_push(&alcf->parts);
    if(part == NULL) {
        return NGX_CONF_ERROR;
    }

    if(alcf->parts.nelts == 1) {
        /* first part */
        part->from = 0;
        part->to = n;
    } else {
        part->from = part[alcf->parts.nelts - 1].to;
        part->to = part->from + n;
    }

    if(part->to > 100) {
        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                            "param too much");

        return NGX_CONF_ERROR;
    }

    part->url.data = value[1].data;
    part->url.len = value[1].len;

    if(cf->args->nelts == 3) {
        if(value[2].len != 7) {
            goto unknown_param;
        }

        if(ngx_strncmp(value[2].data, "default", 7) == 0) {
            alcf->def = alcf->parts.nelts - 1;
        } else {
            goto unknown_param;
        }
    }

    return NGX_CONF_OK;

unknown_param:
    ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                       "unknown param \"%V\"", &value[2]);
    return NGX_CONF_ERROR;
}


static char *ngx_http_abtest_match(ngx_conf_t *cf, ngx_command_t *cmd, void *conf) {
    ngx_str_t                       *value;
    ngx_http_abtest_loc_conf_t      *alcf;
    ngx_http_abtest_match_part_t    *part;

    if(cf->args->nelts < 2 || cf->args->nelts > 3) {
        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                            "invalid param number for abtest");

        return NGX_CONF_ERROR;
    }

    value = cf->args->elts;
    alcf = cf->ctx;

    part = ngx_array_push(&alcf->parts);
    if(part == NULL) {
        return NGX_CONF_ERROR;
    }

    part->pattern = value[0];
    part->url = value[1];

    if(cf->args->nelts == 3) {
        if(ngx_strncmp(value[2].data, "default", 7) == 0) {
            if(alcf->def != 0) {
                ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                                    "default already set");
                return NGX_CONF_ERROR;
            }

            alcf->def = alcf->parts.nelts - 1;
            return NGX_CONF_OK;
        }

        ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                            "invalid param \"%V\" for abtest", &value[2]);

        return NGX_CONF_ERROR;
    }


    return NGX_CONF_OK;
}


static ngx_int_t ngx_http_abtest_postconfiguration(ngx_conf_t *cf) {
    ngx_http_handler_pt                 *h;
    ngx_http_core_main_conf_t           *cmcf;

    cmcf = ngx_http_conf_get_module_main_conf(cf, ngx_http_core_module);
    h = ngx_array_push(&cmcf->phases[NGX_HTTP_REWRITE_PHASE].handlers);
    if(h == NULL) {
        return NGX_ERROR;
    }

    *h = ngx_http_abtest_handler;

    return NGX_OK;
}


static void *ngx_http_abtest_create_loc_conf(ngx_conf_t *cf) {
    ngx_http_abtest_loc_conf_t          *alcf;

    alcf = ngx_pcalloc(cf->pool, sizeof(ngx_http_abtest_loc_conf_t));
    if(alcf == NULL) {
        return NULL;
    }

    return alcf;
}


static char *ngx_http_abtest_percent_init(ngx_conf_t *cf, ngx_http_abtest_loc_conf_t *alcf) {
    ngx_http_abtest_percent_part_t      *parts;
    ngx_uint_t                          i;
    uint32_t                            range;

    alcf->handler = ngx_http_abtest_percent_handler;

    parts = alcf->parts.elts;
    for(i = 0; i < alcf->parts.nelts; ++i) {
        range = (uint32_t)(((float)(parts[i].to - parts[i].from) / (float)100) * NGX_HTTP_ABTEST_SIZE);

        if(i == 0) {
            /* first part */
            parts[i].from = 0;
            parts[i].to = parts->from + range;
        } else {
            parts[i].from = parts[i - 1].to;
            parts[i].to = parts[i].from + range;
        }

    }
    parts[alcf->parts.nelts - 1].to = NGX_HTTP_ABTEST_SIZE;

    return NGX_CONF_OK;
}


static char *ngx_http_abtest_match_init(ngx_conf_t *cf, ngx_http_abtest_loc_conf_t *alcf) {
    ngx_uint_t                          i;
    ngx_http_abtest_match_part_t        *part;
    ngx_regex_compile_t                 rc;
    u_char                              errstr[NGX_MAX_CONF_ERRSTR];
    ngx_str_t                           err;

    err.data = errstr;
    err.len = NGX_MAX_CONF_ERRSTR;
    alcf->handler = ngx_http_abtest_match_handler;

    for(i = 0; i < alcf->parts.nelts; ++i) {
        part = (ngx_http_abtest_match_part_t *)alcf->parts.elts + i;

        ngx_memzero(&rc, sizeof(ngx_regex_compile_t));
        rc.pattern = part->pattern;
        rc.regex = &part->regex;
        rc.pool = cf->pool;
        rc.options = NGX_REGEX_CASELESS;
        rc.err.len = err.len;
        rc.err.data = err.data;

        if(ngx_regex_compile(&rc) != NGX_OK) {
            ngx_conf_log_error(NGX_LOG_EMERG, cf, 0,
                                "regex \"%V\" compile error: %V", (&part->pattern), &err);
            return NGX_CONF_ERROR;
        }

        part->regex.code = rc.regex->code;
        part->regex.extra = rc.regex->extra;
    }

    return NGX_CONF_OK;
}


static ngx_int_t ngx_http_abtest_handler(ngx_http_request_t *r) {
    ngx_http_abtest_loc_conf_t          *alcf;
    ngx_str_t                           url;

    alcf = ngx_http_get_module_loc_conf(r, ngx_http_abtest_module);

    if(alcf->key.len == 0) {
        return NGX_DECLINED;
    }

    if(alcf->handler(r, &url) != NGX_OK) {
        return NGX_HTTP_INTERNAL_SERVER_ERROR;
    }

    ngx_log_debug1(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
                    "abtest internal redirect to \"%V\"", &url);

	if(url.data[0] == '@') {
		return ngx_http_named_location(r, &url);
	}

    return ngx_http_internal_redirect(r, &url, &r->args);
}


static ngx_int_t ngx_http_abtest_match_handler(ngx_http_request_t *r, ngx_str_t *to) {
    ngx_http_abtest_loc_conf_t          *alcf;
    ngx_uint_t                          i;
    ngx_http_abtest_match_part_t        *part;
    ngx_int_t                           n;
    ngx_http_variable_value_t           *value;
    uint32_t                            hash;
    ngx_str_t                           s;

    alcf = ngx_http_get_module_loc_conf(r, ngx_http_abtest_module);
    hash = ngx_murmur_hash2(alcf->key.data, alcf->key.len);

    value = ngx_http_get_variable(r, &alcf->key, hash);
    if(value->not_found) {
        part = (ngx_http_abtest_match_part_t *)alcf->parts.elts + alcf->def;
        ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
                       "no key \"%V\" found in request, use default url \"%V\"",
                       &alcf->key, (&part->url));

        goto done;
    }

    s.data = value->data;
    s.len = value->len;

    for(i = 0; i < alcf->parts.nelts; ++i) {
        part = (ngx_http_abtest_match_part_t *)alcf->parts.elts + i;

        /* ngx_regex_exec is a marco function
         * &part->regex will be &part->regex->code
         * may be a patch is need */
        n = ngx_regex_exec((&part->regex), &s, NULL, 0);

        if(n == NGX_REGEX_NO_MATCHED) {
            ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
                            "\"%V\" do not match \"%V\"", &s, (&part->pattern));

            continue;
        }

        if(n < 0) {
            ngx_log_error(NGX_LOG_ALERT, r->connection->log, 0,
                            "regex faild: %i on \"%V\" using \"%V\"", n, s, (&part->pattern));

            return NGX_ERROR;
        }

        ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
                        "\"%V\" match \"%V\"", &s, (&part->pattern));

        /* match */
        goto done;
    }

    part = (ngx_http_abtest_match_part_t *)alcf->parts.elts + alcf->def;

done:
    to->data = part->url.data;
    to->len = part->url.len;

    return NGX_OK;
}


static ngx_int_t ngx_http_abtest_percent_handler(ngx_http_request_t *r, ngx_str_t *to) {
    ngx_http_abtest_loc_conf_t          *alcf;
    ngx_http_abtest_percent_part_t      *part;
    ngx_http_variable_value_t           *vv;
    uint32_t                            hash, slot;
    ngx_uint_t                          i;

    part = NULL;
    alcf = ngx_http_get_module_loc_conf(r, ngx_http_abtest_module);
    hash = ngx_murmur_hash2(alcf->key.data, alcf->key.len);

    vv = ngx_http_get_variable(r, &alcf->key, hash);
    if(vv->not_found) {
        part = (ngx_http_abtest_percent_part_t *)alcf->parts.elts + alcf->def;
        ngx_log_debug2(NGX_LOG_DEBUG_HTTP, r->connection->log, 0,
                        "no key \"%V\" found in request, use default url \"%V\"",
                        &alcf->key, (&part->url));

        goto done;
    }

    hash = ngx_murmur_hash2(vv->data, vv->len);
    slot = hash & NGX_HTTP_ABTEST_MASK;

    for(i = 0; i < alcf->parts.nelts; ++i) {
        part = (ngx_http_abtest_percent_part_t *)alcf->parts.elts + i;

        if(slot <= part->to) {
            goto done;
        }
    }

done:
    *to = part->url;

    return NGX_OK;
}

